/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "test_common.hpp"
#include <iostream>
#include <vector>
#include <random>
#include <complex>
#include <cmath>
#include <cassert>
#include <chrono>

// ===================== 修复比较函数 (解决重载歧义) =====================
template<typename T>
bool almost_equal(T a, T b, double eps);

template<>
bool almost_equal<float>(float a, float b, double eps) {
    return std::fabs(a - b) < static_cast<float>(eps);
}

template<>
bool almost_equal<double>(double a, double b, double eps) {
    return std::fabs(a - b) < eps;
}

template<typename T>
bool almost_equal(std::complex<T> a, std::complex<T> b, double eps) {
    return almost_equal(a.real(), b.real(), static_cast<T>(eps)) && 
           almost_equal(a.imag(), b.imag(), static_cast<T>(eps));
}

// ===================== 矩阵生成函数 =====================
template<typename T>
void make_lower_csc(std::vector<int>& col_ptr,
                    std::vector<int>& row_ind,
                    std::vector<T>& vals,
                    int n = 8) 
{
    std::mt19937 rng(42);
    std::uniform_real_distribution<double> ur(0.2, 1.0);

    col_ptr.resize(n + 1);
    row_ind.resize(n * (n + 1) / 2);
    vals.resize(row_ind.size());

    int idx = 0;
    for (int j = 0; j < n; ++j) {
        col_ptr[j] = idx;
        for (int r = j; r < n; ++r) {
            row_ind[idx] = r;
            if (r == j)
                vals[idx] = static_cast<T>(ur(rng) + 1.0); // 保证对角远离0
            else
                vals[idx] = static_cast<T>(ur(rng));
            ++idx;
        }
    }
    col_ptr[n] = idx;
}

// ===================== CSC矩阵向量乘法 =====================
template<typename T>
void csc_mv(int n,
            const int* col_ptr, const int* row_ind, const T* vals,
            const T* x, T* y) 
{
    std::fill(y, y + n, T{});
    for (int j = 0; j < n; ++j) {
        for (int p = col_ptr[j]; p < col_ptr[j + 1]; ++p) {
            y[row_ind[p]] += vals[p] * x[j];
        }
    }
}

// ===================== 三角求解器测试模板 =====================
template<typename T, typename SolveFn>
bool run_case(const char* tag,
              SolveFn solve,
              double eps,
              int n = 8) 
{
    /* 1. 构造矩阵 */
    std::vector<int> col_ptr, row_ind;
    std::vector<T> Aval;
    make_lower_csc(col_ptr, row_ind, Aval, n);

    /* 2. 随机生成真实解 */
    std::mt19937 rng(2024);
    std::uniform_real_distribution<double> ur(-1.0, 1.0);
    std::vector<T> x_true(n);
    for (auto& v : x_true) {
        v = static_cast<T>(ur(rng));
    }

    /* 3. 计算右侧向量 b = A·x_true */
    std::vector<T> b(n);
    csc_mv(n, col_ptr.data(), row_ind.data(), Aval.data(),
           x_true.data(), b.data());

    /* 4. 调用三角求解器求解 */
    std::vector<T> x_ref(n);
    auto status = solve(n,
                        col_ptr.data(), row_ind.data(), Aval.data(),
                        x_ref.data(), b.data(),
                        OPENSPBLAS_ZERO_BASE, OPENSPBLAS_LOWER);

    if (status != OPENSPBLAS_STATUS_SUCCESS) {
        std::cout << tag << " : solver returned error " << status << "\n";
        return false;
    }

    /* 5. 比较求解结果 */
    bool success = true;
    for (int i = 0; i < n; ++i) {
        if (!almost_equal(x_true[i], x_ref[i], eps)) {
            if constexpr (std::is_same_v<T, std::complex<float>> || 
                          std::is_same_v<T, std::complex<double>>) {
                std::cout << tag << " : mismatch at i=" << i
                          << "  true=(" << x_true[i].real() << "," << x_true[i].imag() << ")"
                          << "  got=(" << x_ref[i].real() << "," << x_ref[i].imag() << ")\n";
            } else {
                std::cout << tag << " : mismatch at i=" << i
                          << "  true=" << x_true[i] << "  got=" << x_ref[i] << "\n";
            }
            success = false;
        }
    }
    
    if (success) {
        std::cout << tag << " : PASS\n";
    }
    return success;
}

// ===================== 性能测试 =====================
template<typename T>
void test_performance() {
    const int n = 1000;
    const int trials = 10;
    
    std::vector<int> col_ptr, row_ind;
    std::vector<T> Aval;
    make_lower_csc(col_ptr, row_ind, Aval, n);
    
    std::vector<T> b(n), x(n);
    std::mt19937 rng(2024);
    std::uniform_real_distribution<double> ur(-1.0, 1.0);
    for (auto& v : b) v = static_cast<T>(ur(rng));
    
    auto start = std::chrono::high_resolution_clock::now();
    
    for (int i = 0; i < trials; i++) {
        OpenSpB_status status;
        if constexpr (std::is_same_v<T, float>) {
            status = OpenSpB_ssptrsv(n, col_ptr.data(), row_ind.data(), Aval.data(),
                                        x.data(), b.data(), OPENSPBLAS_ZERO_BASE, OPENSPBLAS_LOWER);
        } else if constexpr (std::is_same_v<T, double>) {
            status = OpenSpB_dsptrsv(n, col_ptr.data(), row_ind.data(), Aval.data(),
                                        x.data(), b.data(), OPENSPBLAS_ZERO_BASE, OPENSPBLAS_LOWER);
        } else if constexpr (std::is_same_v<T, std::complex<float>>) {
            status = OpenSpB_csptrsv(n, col_ptr.data(), row_ind.data(), Aval.data(),
                                        reinterpret_cast<void*>(x.data()), 
                                        reinterpret_cast<const void*>(b.data()),
                                        OPENSPBLAS_ZERO_BASE, OPENSPBLAS_LOWER);
        } else {
            status = OpenSpB_zsptrsv(n, col_ptr.data(), row_ind.data(), Aval.data(),
                                        reinterpret_cast<void*>(x.data()), 
                                        reinterpret_cast<const void*>(b.data()),
                                        OPENSPBLAS_ZERO_BASE, OPENSPBLAS_LOWER);
        }
        if (status != OPENSPBLAS_STATUS_SUCCESS) {
            std::cout << "Performance test failed!" << std::endl;
            return;
        }
    }
    
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    std::cout << "Performance for " << typeid(T).name() << " (" << trials << " runs): "
              << duration << " ms" << std::endl;
}

// ===================== 主函数 =====================
TEST(SparseBLASlevel2, spsv) {
    int fails = 0;

    std::cout << "=== Testing Sparse Triangular Solve for CSC Format ===\n";
    
    /* 参数设置 */
    constexpr double EPS_F = 1e-4;
    constexpr double EPS_D = 1e-10;
    constexpr int N = 8;

    /* 测试矩阵 */
    std::cout << "Testing with " << N << "x" << N << " matrix\n";
    
    /* float 测试 */
    fails += !run_case<float>(
        "sspTRSV (float)",
        [](int n, const int* cp, const int* ri, const float* v,
           float* x, const float* b,
           OpenSpB_base_type base, OpenSpB_uplo_type u) {
            return OpenSpB_ssptrsv(n, cp, ri, v, x, b, base, u);
        },
        EPS_F, N);

    /* double 测试 */
    fails += !run_case<double>(
        "dspTRSV (double)",
        [](int n, const int* cp, const int* ri, const double* v,
           double* x, const double* b,
           OpenSpB_base_type base, OpenSpB_uplo_type u) {
            return OpenSpB_dsptrsv(n, cp, ri, v, x, b, base, u);
        },
        EPS_D, N);

    /* complex<float> 测试 */
    fails += !run_case<std::complex<float>>(
        "cspTRSV (cfloat)",
        [](int n, const int* cp, const int* ri, const std::complex<float>* v,
           std::complex<float>* x, const std::complex<float>* b,
           OpenSpB_base_type base, OpenSpB_uplo_type u) {
            return OpenSpB_csptrsv(n, cp, ri, reinterpret_cast<const void*>(v),
                                    reinterpret_cast<void*>(x), 
                                    reinterpret_cast<const void*>(b),
                                    base, u);
        },
        EPS_F, N);

    /* complex<double> 测试 */
    fails += !run_case<std::complex<double>>(
        "zspTRSV (cdouble)",
        [](int n, const int* cp, const int* ri, const std::complex<double>* v,
           std::complex<double>* x, const std::complex<double>* b,
           OpenSpB_base_type base, OpenSpB_uplo_type u) {
            return OpenSpB_zsptrsv(n, cp, ri, reinterpret_cast<const void*>(v),
                                    reinterpret_cast<void*>(x), 
                                    reinterpret_cast<const void*>(b),
                                    base, u);
        },
        EPS_D, N);

    /* 大矩阵测试 */
    const int BIG_N = 500;
    std::cout << "\nTesting with " << BIG_N << "x" << BIG_N << " matrix\n";
    fails += !run_case<std::complex<double>>(
        "zspTRSV-Big (cdouble)",
        [](int n, const int* cp, const int* ri, const std::complex<double>* v,
           std::complex<double>* x, const std::complex<double>* b,
           OpenSpB_base_type base, OpenSpB_uplo_type u) {
            return OpenSpB_zsptrsv(n, cp, ri, reinterpret_cast<const void*>(v),
                                    reinterpret_cast<void*>(x), 
                                    reinterpret_cast<const void*>(b),
                                    base, u);
        },
        EPS_D, BIG_N);

    /* 性能测试 */
    std::cout << "\n=== Performance Testing ===\n";
    test_performance<float>();
    test_performance<double>();
    test_performance<std::complex<float>>();
    test_performance<std::complex<double>>();

    /* 测试总结 */
    std::cout << "\n=== Test Summary ===\n";
    if (fails == 0) {
        std::cout << "All sparse triangular solve tests PASSED!\n";
    } else {
        std::cout << fails << " test(s) FAILED!\n";
    }
}